package com.cy.embeddingdemo.testAgent;

import dev.langchain4j.model.openai.OpenAiChatModel;
import kong.unirest.json.JSONArray;
import kong.unirest.json.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.parquet.format.LogicalTypes;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 高质量的读取pdf内容
 * <p>
 * 1. 数据预处理
 * 文本提取：使用工具（如PyMuPDF、PDFMiner等）从PDF中提取文本。
 * 文本清洗：去除无用的格式符号、页码等，统一文本格式。
 * 2. 知识抽取
 * 实体抽取
 * 基于规则和词典的方法：适用于有明确领域知识和规范的文本，如医学、法律等。优点是简单、准确，缺点是无法处理未登录词。
 * 基于机器学习的方法：如HMM、CRF、BiLSTM-CRF等，适合有一定标注数据的情况。
 * 基于大语言模型的方法：如OpenAI的GPT-3.5-turbo，可以处理复杂的语言模式，但需要精心设计提示词。
 * 关系抽取
 * 监督学习方法：需要大量标注数据，如基于深度学习的模型（BERT、BiLSTM等）。
 * 半监督学习方法：如基于Bootstrap的方法，适合标注数据较少的情况。
 * 大语言模型：如GPT-3.5-turbo，通过提示词引导模型提取关系。
 * 3. 知识融合
 * 实体对齐：使用NLP技术（如句子嵌入模型）计算实体之间的相似度，合并相似实体。
 * 关系整合：对不同来源提取的关系进行整合，消除重复和冲突。
 * 4. 知识表示与存储
 * 知识表示：将抽取的知识转化为RDF三元组或图数据库格式。
 * 存储：选择合适的图数据库（如Neo4j）存储知识图谱。
 * 5. 质量控制
 * 人工校验：对抽取结果进行抽样检查，确保准确性。
 * 持续学习：定期更新知识图谱，加入新的数据和知识。
 * 6. 工具推荐
 * 实体抽取工具：Spacy、jieba（中文）、LAC等。
 * 关系抽取工具：OpenNRE、Stanford OpenIE等。
 * 图数据库：Neo4j、OrientDB等。
 * <p>
 * <p>
 * <p>
 * 具体实施思路：
 * 1 使用 minerU 解析处理pdf
 */
@Slf4j
@RestController
public class Step1 {



    @RequestMapping("attackPath")
    public void doAttack() {
        OpenAiChatModel model = OpenAiChatModel.builder()
                .baseUrl("http://192.168.1.30:8000/v1").
                apiKey("x")
                .timeout( Duration.ofMinutes(60))
                .modelName("Qwen2.5").build();

        String folder = "C:\\数据\\2023";

        extractPDFContents(folder, model);
//        String chat = model.chat("你好，你是谁");
//        System.out.println(chat);
    }


    /**
     * 遍历指定路径下的所有PDF文件，并提取内容存放到列表中
     *
     * @param directoryPath 要遍历的目录路径
     * @param model
     * @return 包含每个PDF文件内容的列表
     */
    public void extractPDFContents(String directoryPath, OpenAiChatModel model) {

        File directory = new File(directoryPath);


        // 遍历目录下的所有文件和子目录
        File[] files = directory.listFiles();
        if (files != null) {
            for (File file : files) {
                // 如果是文件且扩展名是.pdf，则提取内容
                if (file.isFile() && file.getName().toLowerCase().endsWith(".pdf")) {
                    extractPDFContent(file.getAbsolutePath(), model);
                }
                // 如果是目录，则递归遍历
                else if (file.isDirectory()) {
                    extractPDFContents(file.getAbsolutePath(), model);
                }
            }
        }
    }

    /**
     * 提取单个PDF文件的内容
     *
     * @param filePath PDF文件的路径
     * @param model
     * @return PDF文件的内容
     */
    private void extractPDFContent(String filePath, OpenAiChatModel model) {
        try (PDDocument document = PDDocument.load(new File(filePath))) {
            PDFTextStripper stripper = new PDFTextStripper();
            String content = stripper.getText(document);

            // 构造系统提示词
            String systemPrompt = "1、将上传文档的这篇安全报告提炼并总结，把攻击技术部分内容和攻击方式部分内容进行区分\n" +
                    "2、对上传文档的攻击方式部分，通过对攻击方式内容的分析，对攻击方式名称、前提条件、攻击方式、攻击效果、涉及CVE漏洞列表、关联设备、关联APT组织、关联工具等进行总结，以下2个攻击总结为其它报告的总结内容，可以参考如下：\n" +
                    "攻击名称： DDoS攻击\n" +
                    "前提条件： 存在官方网站（如sejm.gov.pl）\n" +
                    "攻击方式： 通过大量请求淹没服务器\n" +
                    "攻击效果： 服务器无法访问\n" +
                    "关联CVE： 无\n" +
                    "关联设备：\n" +
                    "\n" +
                    "攻击名称： 钓鱼攻击\n" +
                    "前提条件： 具备政府工作人员或公众的联系信息（如邮箱）\n" +
                    "攻击方式： 通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问\n" +
                    "攻击效果： 获取用户的凭证信息，可能进一步控制用户终端或进行恶意活动\n" +
                    "关联CVE： 无\n" +
                    "关联设备：\n" +
                    "\n" +
                    "3、对上传文档的攻击方式部分，输出json格式，两个json示例如下：\n" +
                    "{\n" +
                    "  \"AttackSummary\": {\n" +
                    "    \"Attacks\": [\n" +
                    "      {\n" +
                    "        \"AttackName\": \"DDoS攻击\",\n" +
                    "        \"Prerequisites\": [\n" +
                    "          \"存在官方网站（如sejm.gov.pl）\"\n" +
                    "        ],\n" +
                    "        \"AttackMethod\": \"通过大量请求淹没目标服务器\",\n" +
                    "        \"AttackEffect\": \"导致服务器无法访问\",\n" +
                    "        \"Consequences\": [\n" +
                    "          \"破坏目标机构的正常运作\",\n" +
                    "          \"中断公共服务\",\n" +
                    "          \"降低政府公信力\"\n" +
                    "        ]，\n" +
                    "        \"CVE\": [\n" +
                    "          \"\", “” ]，\n" +
                    "      },\n" +
                    "      {\n" +
                    "        \"AttackName\": \"钓鱼攻击\",\n" +
                    "        \"Prerequisites\": [\n" +
                    "          \"具备政府工作人员或公众的联系信息（如邮箱）\"\n" +
                    "        ],\n" +
                    "        \"AttackMethod\": \"通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问\",\n" +
                    "        \"AttackEffect\": [\n" +
                    "          \"获取用户的凭证信息\",\n" +
                    "          \"可能进一步控制用户终端或进行恶意活动\"\n" +
                    "        ],\n" +
                    "        \"Consequences\": [\n" +
                    "          \"造成个人数据泄露\",\n" +
                    "          \"引发财务损失\",\n" +
                    "          \"破坏社会信任\"\n" +
                    "        ]，\n" +
                    "        \"CVE\": [\n" +
                    "          \"\", “” ]，\n" +
                    "      }\n" +
                    "    ],\n" +
                    "    \"AdditionalNotes\": {\n" +
                    "      \"ulnerabilities\": \"报告中未提及具体的溢出漏洞或其他技术细节。\",\n" +
                    "      \"Context\": \"攻击由俄罗斯关联的黑客组织'GhostW riter'实施，动机包括政治报复和破坏关键基础设施。\"\n" +
                    "    }\n" +
                    "  }\n" +
                    "}\n" +
                    "或者这个\n" +
                    "{\n" +
                    "  \"AttackSummary\": {\n" +
                    "    \"Attacks\": [\n" +
                    "      {\n" +
                    "        \"AttackName\": \"Hangul EPS漏洞利用攻击\",\n" +
                    "        \"Prerequisites\": [\n" +
                    "          \"攻击目标为Hangul软件用户\",\n" +
                    "          \"攻击者需要一个包含 CVE-2017-8291 漏洞的 HWP 文件用于初始访问\"\n" +
                    "        ],\n" +
                    "        \"AttackMethod\": \"攻击者将恶意EPS文件嵌入到HWP文档中，通过诱骗目标打开文件触发漏洞，执行恶意代码以下载并执行恶意软件（如M2RAT）\",\n" +
                    "        \"AttackEffect\": [\n" +
                    "          \"窃取目标设备上的敏感信息\",\n" +
                    "          \"包括个人身份信息和通信记录\"\n" +
                    "        ],\n" +
                    "        \"Consequences\": [\n" +
                    "          \"目标用户的个人数据泄露\",\n" +
                    "          \"可能导致身份盗窃或进一步的恶意活动\",\n" +
                    "          \"对攻击目标的机构造成潜在损害\"\n" +
                    "        ],\n" +
                    "        \"CVE\": [\n" +
                    "          \"CVE-2017-8291\"\n" +
                    "        ],\n" +
                    "        \"Devices\": [\n" +
                    "          \"受影响的设备包括移动手机\"\n" +
                    "        ]，\n" +
                    "       \"APTOrganization\": \"APT37\"\n" +
                    "      }\n" +
                    "    ],\n" +
                    "    \"AdditionalNotes\": {\n" +
                    "      \"Vulnerabilities\": \"报告中提到了Hangul EPS漏洞 (CVE-2017-8291) 被用于初始访问。\",\n" +
                    "      \"Context\": \"攻击由RedEyes组织实施，主要针对人權 activist、記者和北韓脫北者，目的是竊取敏感信息。攻击手法包括利用共享記憶體進行命令控制和资料外洩。\"\n" +
                    "    }\n" +
                    "  }\n" +
                    "}\n" +
                    "\n" +
                    "\n" +
                    "\n" +
                    "\n" +
                    "4、对上传文档的攻击技术部分，通过对攻击技术内容的分析，对攻击技术名称、目标系统、实现步骤、攻击效果、涉及CVE漏洞列表、关联设备等进行总结，以下2个攻击总结为其它报告的总结内容，可以参考如下：\n" +
                    "攻击技术名称：内核驱动禁用防御技术（“POOR TRY”驱动技术）\n" +
                    "目标系统：Windows操作系统，特别是目标企业的网络域和企业环境\n" +
                    "实现步骤：1、初始访问：攻击者通过钓鱼邮件或其他方式，投递恶意文档或恶意软件，获得对目标系统的初始访问权限 \n" +
                    "\t2、恶意驱动的部署：攻击者利用恶意DLL（如“POOR TRY”驱动）加载内核模式的驱动程序（如fgt.sys），这些驱动通过Windows服务控制工具（如sc.exe）部署到目标系统 \n" +
                    "\t\t恶意驱动的示例命令：sc create fgt binPath= %TEMP%\\fgt.sys type= kernel sc start fgt\n" +
                    "\t3、禁用安全软件：恶意驱动通过终止目标系统上的抗病毒软件、防火墙或其他安全进程，禁用目标系统的防御机制 \n" +
                    "\t4、横向移动：攻击者利用工具（如SharpHound）进行域信任产品渗透，进一步扩大攻击范围 \n" +
                    "攻击效果：禁用了目标系统的安全防御措施，为后续的恶意活动（如数据窃取、横向移动或勒索软件攻击）创建了有利的环境。\n" +
                    "\t攻击者能够实现持久化（通过恶意驱动和服务）和目标系统的横向移动\n" +
                    "关联CVE： 文档中未提及具体的CVE漏洞\n" +
                    "关联设备：无\n" +
                    "关联工具或软件：\n" +
                    "攻击名称： 钓鱼攻击\n" +
                    "前提条件： 具备政府工作人员或公众的联系信息（如邮箱）\n" +
                    "攻击方式： 通过伪造政府域名和创建虚假网站，利用虚假的财务补偿计划等手段，诱骗员工或公众访问\n" +
                    "攻击效果： 获取用户的凭证信息，可能进一步控制用户终端或进行恶意活动\n" +
                    "关联CVE： 无\n" +
                    "关联设备或软件：Windows；\n" +
                    "关联工具和组件：\t恶意驱动（如fgt.sys）。\n" +
                    "\t\t恶意DLL（如“POOR TRY”驱动）。\n" +
                    "\t\t恶意软件：GuLoader、Cobalt Strike、Netsh、Netwire RAT、Ping、PsExec。\n" +
                    "\t\t渗透测试工具：SharpHound。\n" +
                    "\n" +
                    "5、对上传文档的攻击方式部分，输出json格式，示例如下：\n" +
                    "{\n" +
                    "  \"attack_technology_name\": \"内核驱动禁用防御技术（“POOR TRY”驱动技术）\",\n" +
                    "  \"target_system\": {\n" +
                    "    \"operating_system\": \"Windows\",\n" +
                    "    \"environment\": \"企业网络域和企业环境\"\n" +
                    "  },\n" +
                    "  \"implementation_steps\": [\n" +
                    "    {\n" +
                    "      \"step_number\": 1,\n" +
                    "      \"step_name\": \"初始访问\",\n" +
                    "      \"description\": \"攻击者通过钓鱼邮件或其他方式，投递恶意文档或恶意软件，获得对目标系统的初始访问权限。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"step_number\": 2,\n" +
                    "      \"step_name\": \"恶意驱动部署\",\n" +
                    "      \"description\": \"攻击者利用恶意DLL（如“POOR TRY”驱动）加载内核模式的驱动程序（如`fgt.sys`），并通过Windows服务控制工具（如`sc.exe`）部署到目标系统。例如，用命令`sc create fgt binPath= %TEMP%\\fgt.sys type= kernel`创建服务，然后用`sc start fgt`启动服务。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"step_number\": 3,\n" +
                    "      \"step_name\": \"禁用安全软件\",\n" +
                    "      \"description\": \"恶意驱动终止目标系统上的抗病毒软件、防火墙或其他安全进程，禁用防御机制。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"step_number\": 4,\n" +
                    "      \"step_name\": \"横向移动\",\n" +
                    "      \"description\": \"攻击者利用工具如SharpHound进行域信任产品渗透，扩大攻击范围。\"\n" +
                    "    }\n" +
                    "  ],\n" +
                    "  \"attack_effects\": [\n" +
                    "    {\n" +
                    "      \"effect\": \"禁用目标系统的安全防御措施。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"effect\": \"为后续恶意活动（如数据窃取或勒索软件攻击）创建有利环境。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"effect\": \"实现恶意软件的持久化。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"effect\": \"支持横向移动以扩大攻击范围。\"\n" +
                    "    }\n" +
                    "  ],\n" +
                    "  \"cve_list\": \"无提及具体CVE编号。\",\n" +
                    "  \"associated_devices_or_software\": [\n" +
                    "    {\n" +
                    "      \"device_type\": \"操作系统\",\n" +
                    "      \"name\": \"Windows\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"software_type\": \"恶意驱动\",\n" +
                    "      \"name\": \"fgt.sys, POOR TRY\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"tool\": \"sc.exe\",\n" +
                    "      \"description\": \"用于创建和启动恶意驱动服务。\"\n" +
                    "    },\n" +
                    "    {\n" +
                    "      \"malware\": \"GuLoader, Cobalt Strike, Netsh, Netwire RAT, Ping, PsExec\",\n" +
                    "      \"description\": \"与攻击活动相关的工具和恶意软件。\"\n" +
                    "    }\n" +
                    "  ]\n" +
                    "}\n" +
                    "6、就上面的攻击方式和攻击技术两方面的内容，输出最终的json\n";

// 构造用户提示词
            String userPrompt = "请根据以下文档内容进行回答：\n" + content;

// 将系统提示词和用户提示词组合成对话上下文
            String chatContext = systemPrompt + "\n\n" + userPrompt;
            String aiAnswer = model.chat(chatContext);
            log.info(aiAnswer);

            // 提取 JSON 数据
            String jsonPart = extractJson(aiAnswer);

            // 将 JSON 数据写入文件并追加
            appendJsonToFile(jsonPart, "output.json");

        } catch (IOException e) {
            System.err.println("读取PDF文件时出错: " + filePath);
            e.printStackTrace();
        }
    }

    // 提取 JSON 数据的方法
    public static String extractJson(String text) {
        // 假设 JSON 数据以 "[" 或 "{" 开头，以 "]" 或 "}" 结尾
        int start = text.indexOf('[');
        if (start == -1) {
            start = text.indexOf('{');
        }
        int end = text.lastIndexOf(']');
        if (end == -1) {
            end = text.lastIndexOf('}');
        }

        if (start != -1 && end != -1 && end > start) {
            return text.substring(start, end + 1);
        }
        return null;
    }

    // 将 JSON 数据转换为列表
    public static List<Object> parseJsonToList(String json) {
        List<Object> list = new ArrayList<>();
        try {
            // 如果是 JSON 数组
            JSONArray jsonArray = new JSONArray(json);
            for (int i = 0; i < jsonArray.length(); i++) {
                list.add(jsonArray.get(i));
            }
        } catch (Exception e) {
            // 如果是 JSON 对象
            JSONObject jsonObject = new JSONObject(json);
            list.add(jsonObject);
        }
        return list;
    }

    // 将列表写入文件
    public static void writeListToFile(List<Object> list, String filePath) {
        try (FileWriter file = new FileWriter(filePath)) {
            // 将列表转换为 JSON 格式写入文件
            file.write(list.toString());
            System.out.println("数据已成功写入文件：" + filePath);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 将 JSON 数据写入文件并追加
    public static void appendJsonToFile(String json, String filePath) {
        try (FileWriter file = new FileWriter(filePath, true)) { // true 表示追加模式
            // 写入 JSON 数据
            file.write(json);
            file.write("\n"); // 换行，方便下次追加
            System.out.println("JSON 数据已成功追加到文件：" + filePath);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
